from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes
import base64

class EncryptionError(Exception):
    pass

class DecryptionError(Exception):
    pass

class AESCipher:
    """AES加密实现，支持128/192/256位密钥和ECB/CBC/CTR/OFB/CFB模式。"""
    
    # 添加加密模式常量
    MODE_ECB = AES.MODE_ECB
    MODE_CBC = AES.MODE_CBC
    MODE_CTR = AES.MODE_CTR
    MODE_OFB = AES.MODE_OFB
    MODE_CFB = AES.MODE_CFB
    
    def __init__(self, key, mode=AES.MODE_CBC, iv=None, nonce=None):
        """
        使用密钥和模式初始化AES加密器。
        
        参数:
            key (bytes): 加密密钥（16、24或32字节，分别对应AES-128、AES-192或AES-256）
            mode (int): 加密模式（AES.MODE_CBC、AES.MODE_CTR、AES.MODE_OFB、AES.MODE_CFB或AES.MODE_ECB）
            iv (bytes): 初始化向量（对于CBC、OFB、CFB模式必需）
            nonce (bytes): 随机数（对于CTR模式必需）
        """
        self.key = self._process_key(key)
        self.mode = mode
        
        if mode == AES.MODE_CTR:
            if nonce:
                if isinstance(nonce, str):
                    nonce = base64.b64decode(nonce)
                if len(nonce) != 12:  # CTR模式推荐使用12字节nonce
                    raise ValueError("CTR模式nonce长度必须为12字节")
            else:
                nonce = get_random_bytes(12)
            self.nonce = nonce
            self.iv = None
        elif mode != AES.MODE_ECB:
            if iv:
                if isinstance(iv, str):
                    iv = base64.b64decode(iv)
                if len(iv) != AES.block_size:
                    raise ValueError(f"IV长度必须为{AES.block_size}字节")
            else:
                iv = get_random_bytes(AES.block_size)
            self.iv = iv
            self.nonce = None
        else:
            self.iv = None
            self.nonce = None
    
    def _process_key(self, key):
        """确保密钥满足长度要求。"""
        if not key:
            raise ValueError("密钥不能为空")
            
        if isinstance(key, str):
            key = key.encode('utf-8')
        
        valid_lengths = (16, 24, 32)  # AES-128, AES-192, AES-256
        if len(key) not in valid_lengths:
            key = key.ljust(32, b'\0')[:32]  # 默认使用AES-256
        return key
    
    def encrypt(self, data):
        """
        使用AES加密数据。
        
        参数:
            data: 要加密的字符串或字节
            
        返回:
            tuple: (Base64编码的加密数据, Base64编码的IV/nonce)
            
        异常:
            EncryptionError: 如果加密失败
        """
        try:
            if isinstance(data, str):
                data = data.encode('utf-8')
            
            if self.mode == AES.MODE_ECB:
                cipher = AES.new(self.key, self.mode)
                iv = None
            elif self.mode == AES.MODE_CTR:
                cipher = AES.new(self.key, self.mode, nonce=self.nonce)
                iv = self.nonce
            else:
                cipher = AES.new(self.key, self.mode, self.iv)
                iv = self.iv
            
            if self.mode == AES.MODE_CTR:
                encrypted_data = cipher.encrypt(data)  # CTR模式不需要填充
            else:
                padded_data = pad(data, AES.block_size)
                encrypted_data = cipher.encrypt(padded_data)
            
            result = base64.b64encode(encrypted_data).decode('utf-8')
            iv_str = base64.b64encode(iv).decode('utf-8') if iv else None
            
            return result, iv_str
            
        except Exception as e:
            raise EncryptionError(f"AES加密失败: {str(e)}")

    def decrypt(self, encrypted_data, iv=None):
        """
        解密AES加密的数据。
        
        参数:
            encrypted_data: Base64编码的加密数据
            iv: Base64编码的IV/nonce（对于CBC、CTR、OFB、CFB模式必需）
            
        返回:
            解密后的字符串
            
        异常:
            DecryptionError: 如果解密失败
        """
        try:
            encrypted_bytes = base64.b64decode(encrypted_data)
            
            if self.mode == AES.MODE_ECB:
                cipher = AES.new(self.key, self.mode)
            elif self.mode == AES.MODE_CTR:
                if not iv:
                    raise ValueError("CTR模式需要提供nonce")
                nonce_bytes = base64.b64decode(iv)
                cipher = AES.new(self.key, self.mode, nonce=nonce_bytes)
            else:
                if not iv:
                    raise ValueError("非ECB模式需要提供IV")
                iv_bytes = base64.b64decode(iv)
                cipher = AES.new(self.key, self.mode, iv_bytes)
            
            if self.mode == AES.MODE_CTR:
                decrypted_data = cipher.decrypt(encrypted_bytes)  # CTR模式不需要解填充
            else:
                decrypted_data = unpad(cipher.decrypt(encrypted_bytes), AES.block_size)
            
            return decrypted_data.decode('utf-8')
            
        except Exception as e:
            raise DecryptionError(f"AES解密失败: {str(e)}")